import jsonlines

from datasets import load_dataset
from tqdm import tqdm
#raw_dataset = load_dataset("../Infinity_Instruct/7M",split="train")
import json

raw_dataset = []
with jsonlines.open("/share/project/lijijie/tools/instruction_follow/0625/7M/7M_0712_math_plus_system_release_0802.jsonl") as f:
    for line in tqdm(f):
        raw_dataset.append(line)

import sys
idx = sys.argv[1]
contexts = {}

finish_dataset = []
for i in range(5):
    with jsonlines.open(f"7M_label_{i}.jsonl") as f:
        for line in f:
            #if line["conversations"][-1]["value"] not in contexts:
            if len(line["label"]["ability"])==0:
                continue
            #finish_dataset.append(line)
            contexts[line["conversations"][-2]["value"][:100] + line["conversations"][-1]["value"][:100]] = line
            #contexts.append(" ".join([a["value"] for a in line["conversations"]]))
        
with jsonlines.open(f"7M_label_initial.jsonl") as f:
    for line in f:
        #if line["conversations"][-1]["value"] not in contexts:
        if len(line["label"]["ability"])==0:
            continue
        #finish_dataset.append(line)
        contexts[line["conversations"][-2]["value"][:100] + line["conversations"][-1]["value"][:100]] = line
#for line in whole_dataset:
#    raw_dataset.append()

#with jsonlines.open(f"") as f:
#    for line in f:
#        raw_dataset.append(line)

#for d in tqdm(raw_dataset):
#    contexts.add(d["conversations"][-1]["value"][:100])


#with jsonlines.open(f"7M_label_{idx}.jsonl") as f:
#    for line in tqdm(f):
#        if len(line["label"]["ability"])==0:
#            continue
#        contexts.add(str(line["id"])+line["conversations"][-1]["value"])
        #if line["chosen"][0]["content"] not in contexts:
        #    datas.append(line)
datas = []
for d in tqdm(raw_dataset):
    if (d["conversations"][-2]["value"][:100] + d["conversations"][-1]["value"][:100]) in contexts:
        d["label"] = contexts[d["conversations"][-2]["value"][:100] + d["conversations"][-1]["value"][:100]]["label"]
    else:
        #import pdb
        #pdb.set_trace()
        datas.append(d)
    finish_dataset.append(d)

with jsonlines.open(f"7M_label_filter_whole.jsonl","w") as wf:
    for line in datas:
        wf.write(line)


with jsonlines.open(f"7M_label_tagged.jsonl","w") as wf:
    for line in finish_dataset:
        wf.write(line)
